# -*- coding: UTF-8 -*-
# *******************************************************************
# File Name: test
# > Author: 04000387
# > Created Time: 2024/12/23 16:34
# *******************************************************************

# from datas import ReadTxt
#
# import cv2
# from pathlib import Path
# import matplotlib.pyplot as plt
# from torchvision.transforms import v2
# import torch

# transformers = v2.Compose([
#     v2.ToImage(),
#     v2.ToDtype(torch.float32, scale=True),
#     v2.Resize(size=(224, 224)),
#     v2.Normalize(mean=[0.48235, 0.45882, 0.40784], std=[1.0 / 255, 1.0 / 255, 1.0 / 255])
# ])

# transformers = v2.Compose([
#     v2.ToImage(),
#     v2.ToDtype(torch.float32, scale=True),
#     v2.Resize(size=(600, 800))
# ])
#
# path = "E:\\01_git\\generate_dir\\generate\\img\\0000a003-cf18-42ff-86e4-86b63c613f2d.png"
#
# path = "E:\\01_git\\generate_dir\\test\\train_full_images_0\\gt_148.jpg"
#
# img_path = path
# img = cv2.imread(img_path)
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
#
#
# transformer_img = transformers(img)
# print(transformer_img.size())
#
# transformer_img = transformer_img.permute(1, 2, 0)
#
# plt.imshow(transformer_img)
# plt.show()


from utils import plot_box
from datas import ReadWithBoxData, WithBoxCollectFn
from torch.utils.data import DataLoader
from transformers import BertTokenizer
import torch
from torchvision.transforms import v2

transformers = v2.Compose([
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Resize(size=(224, 224))
])


model_path = "D:\\models\\ernie"
tokenizer = BertTokenizer.from_pretrained(model_path)

path = "../generate/pure"
reader = ReadWithBoxData(path)

collate_fn = WithBoxCollectFn(tokenizer, transformers)

loader = DataLoader(reader, 8, True, collate_fn=collate_fn)
for img, text, boxes in loader:
    plot_box(img[0].permute(1, 2, 0), boxes[0] * 224.0)
    break

# from tool import batch_generate_pure_background_img
# text_path = "D:\\deepl_git\\generate_dir\\txt"
# save_path = "../generate/pure"
# num = 64
# batch_generate_pure_background_img(text_path, save_path, _num=num)


